﻿#pragma once
#include <iostream>
#include <cassert>
using namespace std;

enum COLOUR
{
	RED,
	BLACK
};

template<class K, class V>
struct RBTNode
{
	pair<K, V> _kv;
	RBTNode<K,V>* _left = nullptr;
	RBTNode<K,V>* _right = nullptr;
	RBTNode<K,V>* _parent = nullptr;
	COLOUR _color;

	RBTNode(const K& key = K(), const V& value = V())
		:_kv(key, value)
	{}
	RBTNode(const pair<K, V>& kv)
		:_kv(kv)
	{}
};

template<class K, class V>
class RBTree
{
public:
	typedef RBTNode<K, V> Node;

	//右单旋
	void RotateR(Node* son)
	{
		Node* parent = son->_parent;
		Node* sonR = son->_right;
		Node* gparent = parent->_parent;
		son->_right = parent;
		parent->_parent = son;
		parent->_left = sonR;
		if (sonR)
			sonR->_parent = parent;

		son->_parent = gparent;
		if (gparent == nullptr)
		{
			//说明parent为根结点，旋转后，son为根结点
			_root = son;
		}
		else
		{
			if (gparent->_left == parent)
				gparent->_left = son;
			else if (gparent->_right == parent)
				gparent->_right = son;
			else
				assert(false);
		}
	}

	//左单旋
	void RotateL(Node* son)
	{
		Node* parent = son->_parent;
		Node* gparent = parent->_parent;
		Node* sonL = son->_left;

		son->_left = parent;
		parent->_parent = son;
		parent->_right = sonL;
		if (sonL)
			sonL->_parent = parent;

		son->_parent = gparent;
		if (gparent == nullptr)
		{
			_root = son;
		}
		else
		{
			if (gparent->_left == parent)
			{
				gparent->_left = son;
			}
			else if (gparent->_right == parent)
			{
				gparent->_right = son;
			}
			else
				assert(false);
		}
	}
	void Insert(const pair<K, V>& kv)
	{
		Node* ptr = new Node(kv);
		ptr->_color = RED;
		if (_root == nullptr)
		{
			_root = ptr;
			_root->_color = BLACK;
		}
		else
		{
			//遍历整棵二叉搜索树，找到应处位置
			Node* cur = _root;
			Node* pre = nullptr;
			while (cur)
			{
				pre = cur;
				if (ptr->_kv.first < cur->_kv.first)
				{
					cur = cur->_left;
				}
				else if (ptr->_kv.first > cur->_kv.first)
				{
					cur = cur->_right;
				}
				else//暂且不处理相等的情况
				{
					delete ptr;
					return;
				}

			}
			ptr->_parent = pre;
			if (ptr->_kv.first < pre->_kv.first)
			{
				pre->_left = ptr;
			}
			else
			{
				pre->_right = ptr;
			}
		}

		Node* cur = ptr;
		Node* parent = ptr->_parent;
		//插入结点，初始化完成，开始变色+旋转

		//1.如果parent为空(根结点插入情况)或者parent为黑，不需要处理

		//2. 需要处理的是parent不为空并且parent为红的情况
		while (parent && parent->_color == RED)
		{
			Node* grandParent = parent->_parent;
			if (parent == grandParent->_left)
			{
				Node* uncle = grandParent->_right;
				//分三种情况，但两种处理来讨论

				//1.叔叔存在且为红，不断向上调整,此时仅变色即可
				//注意，不能写成循环的形式
				if (uncle && uncle->_color == RED)
				{
					parent->_color = uncle->_color = BLACK;
					grandParent->_color = RED;

					cur = grandParent;
					parent = cur->_parent;
				}

				//2.叔叔存在且为黑，或叔叔不存在，此时需变色+旋转
				//同样内部不能写成循环的形式
				else if (uncle == nullptr || uncle && uncle->_color == BLACK)
				{
					if (cur == parent->_left)
					{
						//单旋+变色
						RotateR(parent);

						parent->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else if (cur == parent->_right)
					{
						//双旋+变色
						RotateL(cur);
						RotateR(cur);

						cur->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else
					{
						assert(false);
					}
				}
				else
				{
					assert(false);
				}

			}
			else if (parent == grandParent->_right)
			{
				Node* uncle = grandParent->_left;
				//分三种情况，但两种处理来讨论

				//1.叔叔存在且为红，不断向上调整,此时仅变色即可
				//注意，不能写成循环的形式
				if (uncle && uncle->_color == RED)
				{
					parent->_color = uncle->_color = BLACK;
					grandParent->_color = RED;

					cur = grandParent;
					parent = cur->_parent;
				}

				//2.叔叔存在且为黑，或叔叔不存在，此时需变色+旋转
				//同样内部不能写成循环的形式
				else if (uncle == nullptr || uncle && uncle->_color == BLACK)
				{
					if (cur == parent->_right)
					{
						//单旋+变色
						RotateL(parent);

						parent->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else if (cur == parent->_left)
					{
						//双旋+变色
						RotateR(cur);
						RotateL(cur);

						cur->_color = BLACK;
						grandParent->_color = RED;
						break;
					}
					else
					{
						assert(false);
					}
				}
				else
				{
					assert(false);
				}
			}
			else
			{
				assert(false);
			}
		}
		_root->_color = BLACK;//无脑将根结点的颜色保证为黑色
	}

	bool Check(Node* root, int blackNum, const int refNum)
	{
		if (root == nullptr)
		{
			// 前序遍历⾛到空时，意味着⼀条路径⾛完了
			//cout << blackNum << endl;
			if (refNum != blackNum)
			{
				cout << "存在黑色结点的数量不相等的路径" << endl;
				return false;
			} 
			return true;
		} 
		// 检查孩⼦不太⽅便，因为孩⼦有两个，且不⼀定存在，反过来检查⽗亲就⽅便多了
			if (root->_color == RED && root->_parent->_color == RED)
			{
				cout << root->_kv.first << "存在连续的红色结点" << endl;
				return false;
			} 
			if(root->_color == BLACK)
			{
				blackNum++;
			}
	
			return Check(root->_left, blackNum, refNum) && Check(root->_right, blackNum, refNum);
	} 
	bool IsBalance()
	{
		if (_root == nullptr)
			return true;
		if (_root->_color == RED)
			return false;
		// 参考值
		int refNum = 0;
		Node* cur = _root;
		while (cur)
		{
			if (cur->_color == BLACK)
			{
				++refNum;
			} 
			cur = cur->_left;
		} 
		return Check(_root, 0, refNum);
	}
private:
	Node* _root = nullptr;//根结点
};


